/*
Copyright (C) 2022-2024 Inspur Corp.

SPDX-License-Identifier: Apache-2.0
*/
package gm

import (
	"crypto/rand"
	"encoding/hex"
	"github.com/tjfoc/gmsm/sm2"
	"icep/common/logger"
	"math/big"
	"strings"
)

var log = logger.GetLogger("common.crypto")

func GenerateSM2Keys() (string, string, error) {
	privateKey, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		return "", "", err
	}
	dHex := hex.EncodeToString(privateKey.D.Bytes())
	publicKey := privateKey.PublicKey
	pubBytes := sm2.Compress(&publicKey)
	pubKeyHex := hex.EncodeToString(pubBytes)
	prvKeyHex := strJoin(dHex, pubKeyHex)
	return prvKeyHex, pubKeyHex, nil
}

func GetPubKey(pubKeyHex string) *sm2.PublicKey {
	pubKeyBytes, err := hex.DecodeString(pubKeyHex)
	if err != nil {
		return nil
	}
	pubKey := sm2.Decompress(pubKeyBytes)
	return pubKey
}

func GetPrvKey(prvKeyHex string) *sm2.PrivateKey {
	dHex := prvKeyHex[:64]
	pubKeyHex := prvKeyHex[64:]
	pubKey := GetPubKey(pubKeyHex)
	prvKeyBytes, err := hex.DecodeString(dHex)
	if err != nil {
		return nil
	}
	d := new(big.Int)
	d.SetBytes(prvKeyBytes)
	prvKey := &sm2.PrivateKey{
		PublicKey: *pubKey,
		D:         d,
	}
	return prvKey
}

func Sign(prvKeyHex string, message []byte) string {
	prvKey := GetPrvKey(prvKeyHex)
	signBytes, err := prvKey.Sign(rand.Reader, message, nil)
	if err != nil {
		log.Errorf("sm2 sign error: %s", err)
		return ""
	}
	sign := hex.EncodeToString(signBytes)
	return sign
}

func Verify(pubKeyHex string, message []byte, sign string) bool {
	pubKey := GetPubKey(pubKeyHex)
	bytes, err := hex.DecodeString(sign)
	if err != nil {
		log.Errorf("decode signature error: %v", err)
		return false
	}
	result := pubKey.Verify(message, bytes)
	return result
}

func EncryptByPubKey(pubKeyHex string, message []byte) string {
	pubKey := GetPubKey(pubKeyHex)
	asn1, err := pubKey.EncryptAsn1(message, rand.Reader)
	if err != nil {
		log.Errorf("encrypt by public key error: %v", err)
		return ""
	}
	result := hex.EncodeToString(asn1)
	return result
}

func DecryptByPrvKey(prvKeyHex string, encryptedData string) []byte {
	prvKey := GetPrvKey(prvKeyHex)
	bytes, err := hex.DecodeString(encryptedData)
	if err != nil {
		log.Errorf("decode encrypted data error: %v", err)
		return nil
	}
	result, err := prvKey.DecryptAsn1(bytes)
	if err != nil {
		log.Errorf("decrypt by private key error: %v", err)
		return nil
	}
	return result
}

func strJoin(strs ...string) string {
	var builder strings.Builder
	for _, str1 := range strs {
		builder.WriteString(str1)
	}
	return builder.String()
}
